Purpose

Overview

DFlash works much like P-EAGLE (see #32887), but with a major architectural change: it uses bidirectional attention between the query tokens (the last sampled token from the base model plus a bunch of placeholder mask tokens) and the context states, which are the target model's hidden states from the prefill or accepted tokens.

To implement this, I introduce an extra operation that lives outside of the main model execution, which populates the KV cache with the context states directly. Though not exposed to the standard set of torch.compile and CUDA graph optimizations, handling the context in this way allows us to use async scheduling (by writing the full set of context states to the cache and then using seq_lens_gpu to ignore the rejected ones), as well as enabling (piecewise) CUDA graphs for the main forward pass over the query tokens.

(Solved) Because the attention is bidirectional and structured in this way, we cannot allow rejected tokens to stay in the batch as their states would corrupt the rest of the calculation, unlike in standard causal attention where we can simply omit them and sample from an earlier position. Therefore, "disable_padded_drafter_batch" is required and will be enabled when using DFlash, disabling async scheduling as a consequence. At this time I cannot think of a good way around this problem.

Additionally, a small selection of kernel backends actually support non-causal attention, and a smaller set additionally include support for gpt-oss style "sinks". Flash Attention with Qwen3-8B is used as a test for this implementation, as Triton Attention and FlashInfer (TRTLLM) attention both do not support non-causality. A follow-up work here would be to allow different attention backends for the drafter and the target model. This is out-of-scope for this PR, but would enable a broader set of compatibility for DFlash models.

(Solved) Finally, this new architecture requires extra logic to handle the new input shapes and attention metadata. It is not as simple as P-EAGLE in which we can share code with other modes of speculation, since the sizes and contents of the input tensors are all different. One compatibility component in which I am not absolutely confident is the CUDA graph support: even in piecewise mode, I am not sure how the "padded" query tokens interact with the unpadded context: is it safe to have operations on the context slice of Q inside the attention op, or must it be moved into a custom op to avoid issues? It has not yet been functional to enable torch.compile for the DFlash drafter, likely for this reason.

Implementation Details

The model is implemented in qwen3_dflash.py and the core speculation logic is in dflash.py. Similarly to draft model speculative decoding, DFlash has a unique speculation paradigm that requires some refactoring of eagle.py in order to support cleanly.

Specifically, build_model_inputs_first_pass and build_per_layer_attn_metadata have been introduced in eagle.py so that dflash.py can override them.

Initial implementations inlined the logic for DFlash into eagle.py, but similarly to #24322 the added branching and complexity would (in my opinion) lead to a fairly cluttered EAGLE implementation. In this PR I have tried to separate the concerns to a reasonable degree, so that maintenance of the existing EAGLE pathway is not burdened.

Testing

I add unit tests for both correctness (GSM8k) and acceptance rate checks for Qwen3-8B with DFlash. I am able to reproduce almost exactly the acceptance rate values reported in the DFlash paper, and have included this in the test suite with an exact reproduction of their evaluation setup. This should be a valuable resource and an easily extensible tool to new DFlash models as they continue to evolve and expand the DFlash model family.

In local tests on 1xB200, both Qwen3-8B and Qwen3.5-9B pass the test suite using both FA2 and FA4 on 1xB200.

Usage

vllm serve Qwen/Qwen3-8B \
    --speculative-config '{"method": "dflash", "model": "z-lab/Qwen3-8B-DFlash-b16", "num_speculative_tokens": 15}' \
    --attention-backend flash_attn \
    --max-num-batched-tokens 32768

Benchmarking

The latest DFlash implementation on this branch is optimized for low-latency performance. I measured the following speedups using the DFlash methodology (https://github.com/z-lab/dflash)

Datasets covered are Alpaca, GSM8k, HumanEval, Math500, MBPP, MT-Bench. Number of spec tokens for DFlash is 15.

Alpaca

Output tok/s
conc  baseline    dflash  speedup
----  --------  --------  -------
  32  5,268.64  6,210.65    1.179
  16  3,079.46  5,015.65    1.629
   8  1,581.42  3,339.16    2.111
   1    209.46    590.48    2.819

GSM8k

Output tok/s
conc  baseline    dflash  speedup
----  --------  --------  -------
  32  5,261.70  8,433.19    1.603
  16  3,079.63  6,534.04    2.122
   8  1,596.48  4,332.12    2.714
   1    211.35    741.56    3.509

HumanEval

Output tok/s
conc  baseline     dflash  speedup
----  --------  ---------  -------
  32  5,002.62  10,257.62    2.050
  16  3,043.47   8,131.09    2.672
   8  1,583.48   5,433.65    3.431
   1    209.16     969.07    4.633

Math500

Output tok/s
conc  baseline    dflash  speedup
----  --------  --------  -------
  32  5,267.27  9,836.42    1.867
  16  3,087.45  7,668.98    2.484
   8  1,598.57  5,096.36    3.188
   1    211.31    849.22    4.019

MBPP

Output tok/s
conc  baseline    dflash  speedup
----  --------  --------  -------
  32  5,120.32  8,289.34    1.619
  16  3,099.78  6,677.64    2.154
   8  1,584.55  4,487.97    2.832
   1    209.56    810.25    3.866

MT-Bench

Output tok/s
conc  baseline    dflash  speedup
----  --------  --------  -------
  32  5,217.68  6,787.14    1.301
  16  3,082.95  5,297.45    1.718
   8  1,591.69  3,536.00    2.222
   1    210.29    627.01    2.982